"""数据生成器演示脚本

此脚本展示如何使用DataGenerator类生成不同类型的测试数据，包括用户数据、订单数据、交易数据和自定义数据。
"""

import os
import random
from utils.data_generator import DataGenerator
from faker import Faker
# 确保数据目录存在
data_dir = './data/generated_data'
os.makedirs(data_dir, exist_ok=True)

def generate_small_demo_data():
    """生成小批量演示数据"""
    print("=" * 80)
    print("生成小批量演示数据")
    print("=" * 80)
    
    generator = DataGenerator()
    
    # # 1. 生成100条用户数据
    # users_df = generator.generate_user_data(100, f'{data_dir}/users_small.csv')
    # print(f"生成了{len(users_df)}条用户数据")
    # print("用户数据样例:")
    # print(users_df.head())
    # print()

    # # 2. 生成1000条订单数据（使用上面生成的用户ID）
    # user_ids = users_df['user_id'].tolist()
    # orders_df = generator.generate_order_data(1000, user_ids, f'{data_dir}/orders_small.csv')
    # print(f"生成了{len(orders_df)}条订单数据")
    # print("订单数据样例:")
    # print(orders_df.head())
    # print()
    
    # 3. 生成500条交易数据
    transactions_df = generator.generate_transaction_data(500, f'{data_dir}/transactions_small.csv')
    print(f"生成了{len(transactions_df)}条交易数据")
    print("交易数据样例:")
    print(transactions_df.head())
    print()
    
    # 4. 生成自定义数据
    custom_schema = {
        'id': lambda faker, i: i + 1,
        'name': lambda faker, i: faker.name(),
        'age': {'type': 'int', 'min': 18, 'max': 65},
        'department': {'type': 'choice', 'choices': ['销售', '研发', '人事', '财务', '市场']},
        'salary': {'type': 'float', 'min': 3000, 'max': 20000},
        'hire_date': {'type': 'date', 'start_date': '-5y', 'end_date': 'now'},
        'active': True
    }
    employees_df = generator.generate_custom_data(200, custom_schema, f'{data_dir}/employees_small.csv')
    print(f"生成了{len(employees_df)}条自定义数据")
    print("自定义数据样例:")
    print(employees_df.head())
    print()

def generate_large_data_demo():
    """生成大量数据演示"""
    print("=" * 80)
    print("生成大量数据演示")
    print("=" * 80)
    print("注意：生成大量数据可能需要较长时间和较多内存")

    #  生成1000万条订单数据
    # 需要先有用户数据
    # user_ids = list(range(1, 100001))  # 假设有10万个用户
    # generator.generate_order_data(10000000, user_ids, f'{data_dir}/orders_large.parquet')

def generate_custom_schema_demo():
    """自定义数据模式演示"""
    print("=" * 80)
    print("自定义数据模式演示")
    print("=" * 80)

    generator = DataGenerator()

    # 示例1：生成添加支付数据
    product_schema = {
        'productCode': lambda faker, i: f'{random.randint(10000000, 99999999):08d}',
        'productName': lambda faker, i: faker.name()[:2]  # 截取前两个汉字
    }

    products_df = generator.generate_custom_data(500, product_schema, f'{data_dir}/products_custom.csv')
    print("商品数据样例:")
    print(products_df.head())
    print()


if __name__ == "__main__":
    print("数据生成器演示开始")
    print()
    
    # 生成小批量演示数据
    # generate_small_demo_data()
    
    # 生成自定义模式数据演示
    generate_custom_schema_demo()
    
    # 生成大量数据演示（默认注释，需要时取消注释）
    # generate_large_data_demo()
    
    print("数据生成器演示结束")
    print(f"所有生成的数据都保存在 {data_dir} 目录下")